import os
import numpy as np
import matplotlib.pyplot as plt


def ideal(x, E, yield_stress):
    ret = 0
    if E * x < yield_stress:
        ret = E * x
    else:
        ret = yield_stress
    return ret

def bilinear(x, E, Et, yield_stress):
    ret = 0
    if E * x < yield_stress:
        ret = E * x
    else:
        ret = yield_stress + Et * (x - yield_stress / E)
    return ret

def unipowr(x, E, n, yield_stress):
    ret = 0
    k = yield_stress / ((yield_stress / E) ** n)
    if E * x < yield_stress:
        ret = E * x
    else:
        ret = k * x ** n
    return ret


if __name__ == "__main__":
    # importing openVFIFE data
    root = "/media/tan/_dde_data/projects/ginkgo_vfife/tests/materials"

    E = 2.06E11
    Et = 2.06E9
    yield_stress = 235E6

    # compare UniIdeal
    fname = os.path.join(root, "uniideal.csv")
    data = np.loadtxt(fname, delimiter=",")

    x = np.arange(0, 0.01, 1e-6)
    y = np.array([ideal(i, E, yield_stress) for i in x])

    fig, ax = plt.subplots()
    ind = 1000
    ax.plot(data[:,0], data[:,1], c="black", lw=2, label="openVFIFE")
    ax.plot(x, y, c="red", lw=2, label="Theory")
    xmin, xmax = np.min(data[:,0]), np.max(data[:,0])
    ax.hlines(yield_stress, xmin=xmin, xmax=xmax, color="blue", lw=1,
        label="upper")
    ax.hlines(-yield_stress, xmin=xmin, xmax=xmax, color="blue", lw=1.5,
        label="lower")
    plt.legend(loc="best")
    plt.savefig("UniIdeal.svg")
    # plt.show()
    plt.close(fig)

    # compare UniIdeal
    fname = os.path.join(root, "unibilinear0.0.csv")
    data1 = np.loadtxt(fname, delimiter=",")

    fname = os.path.join(root, "unibilinear0.5.csv")
    data2 = np.loadtxt(fname, delimiter=",")

    fname = os.path.join(root, "unibilinear1.0.csv")
    data3 = np.loadtxt(fname, delimiter=",")

    y = np.array([bilinear(i, E, Et, yield_stress) for i in x])
    # y = np.array([unipowr(i, E, 0.5, yield_stress) for i in x])
    fig, ax = plt.subplots()
    ind = 1000
    ax.plot(data1[:,0], data1[:,1], c="black", lw=2, label="m=0.0, Kinematic")
    ax.plot(data2[:,0], data2[:,1], c="green", lw=2, label="m=0.5, Mixed")
    ax.plot(data3[:,0], data3[:,1], c="yellow", lw=2, label="m=1.0, Isotropic")
    ax.plot(x, y, c="red", lw=2, label="Theory")
    xmin, xmax = np.min(data[:,0]), np.max(data[:,0])
    ax.hlines(yield_stress, xmin=xmin, xmax=xmax, color="blue", lw=1,
        label="upper")
    ax.hlines(-yield_stress, xmin=xmin, xmax=xmax, color="blue", lw=1.5,
        label="lower")
    plt.legend(loc="best")
    plt.savefig("UniBilinear.svg")
    # plt.show()
    plt.close(fig)



